﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Xml.Linq;
using System.Xml.XPath;
using System.Diagnostics;
using nVentive.Umbrella.Extensions;
using System.Runtime.InteropServices;
using System.Text.RegularExpressions;
using System.Management;
using System.Xml;
using nVentive.Umbrella;

namespace HyperVMover
{
    public class HyperVUtils
    {
        private readonly static string _defaultHyperVRootPath = Path.Combine(
                    Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData),
                    @"Microsoft\Windows\Hyper-V");

        private static bool enableWrite = true;

        public static event Action<string> LogMessage = delegate { };

        [DllImport("kernel32.dll", EntryPoint = "CreateSymbolicLink")]
        static extern bool _CreateSymbolicLink(string lpSymlinkFileName, string lpTargetFileName, int dwFlags);

        public static void AttachVM(string path)
        {
            AttachVM(null, path, null);
        }

        public static void AttachVM(string hyperVRootPath, string path, string hostMachine)
        {
            if (hyperVRootPath == null)
            {
                hyperVRootPath = _defaultHyperVRootPath;
            }

            FormatLogMessage("Importing [{0}] using [{1}] on [{2}]", path, hyperVRootPath, hostMachine != null ? hostMachine : ".");

            var vmPath = Path.Combine(path, "Virtual Machines");
            var snapshotsPath = Path.Combine(path, "Snapshots");

            if (Directory.Exists(ConvertToUNC(vmPath, hostMachine)))
            {
                var vmFiles = Directory.GetFiles(ConvertToUNC(vmPath, hostMachine), "*.xml");

                if (vmFiles.Length == 1)
                {
                    var configurationFile = vmFiles[0];
                    var configurationFileSymlink = hyperVRootPath + "\\Virtual Machines\\" + Path.GetFileName(configurationFile);

                    var originalBasePath = GuessOriginalBasePath(configurationFile);

                    if (originalBasePath != null)
                    {
                        FormatLogMessage("Guessed original path as [{0}]", originalBasePath);

                        var snapshotsFiles = from sourceSnapshotFile in GetFiles(ConvertToUNC(snapshotsPath, hostMachine), "*.xml")
                                             select new
                                             {
                                                 Source = sourceSnapshotFile,
                                                 Symlink = hyperVRootPath + "\\Snapshots\\" + Path.GetFileName(sourceSnapshotFile)
                                             };

                        // Update paths
                        UpdatePaths(configurationFile, originalBasePath, path);
                        snapshotsFiles.ForEach(f => UpdatePaths(f.Source, originalBasePath, path));

                        // Create links
                        FormatLogMessage("Creating new symlinks to snapshot files...");
                        snapshotsFiles.ForEach(f => CreateSymbolicLink(ConvertToUNC(f.Symlink, hostMachine), ConvertToLocal(f.Source, hostMachine), 0));

                        // Set the ACLs
                        string[] aclUserName = new[]{
                            "NT VIRTUAL MACHINE\\" + Path.GetFileNameWithoutExtension(configurationFile),
                            "NT AUTHORITY\\NETWORK SERVICE"
                        };

                        // Set ACLs for snapshot symlinks
                        FormatLogMessage("Setting ACLs snapshot symlink files...");
                        snapshotsFiles.ForEach(f => SetAcl(ConvertToUNC(f.Symlink, hostMachine), aclUserName, true, hostMachine));

                        // Set the ACLs for all files in actual VM directories
                        FormatLogMessage("Setting ACLs snapshot for all virtual machine files...");
                        SetAcl(ConvertToUNC(path, hostMachine), aclUserName, false, hostMachine);

                        // Now create the symlink to the configuration file
                        FormatLogMessage("Creating new symlinks to VM configuration files...");
                        CreateSymbolicLink(ConvertToUNC(configurationFileSymlink, hostMachine), ConvertToLocal(configurationFile, hostMachine), 0);

                        // And set the ACL for that link.
                        FormatLogMessage("Setting ACLs snapshot VM configuration file...");
                        SetAcl(ConvertToUNC(configurationFileSymlink, hostMachine), aclUserName, true, hostMachine);
                    }
                    else
                    {
                        throw new Exception("Cannot guess directories to update (At least one drive is needed)");
                    }
                }
                else
                {
                    throw new Exception("Too many VMs in the specified directory.");
                }
            }
            else
            {
                throw new Exception("No VM can be found in the specified directory.");
            }
        }

        private static string[] GetFiles(string path, string filter)
        {
            if (Directory.Exists(path))
            {
                return Directory.GetFiles(path, filter);
            }

            return new string[0];
        }

        private static string ConvertToUNC(string path, string hostMachine)
        {
            if (string.IsNullOrEmpty(hostMachine))
                return path;

            return string.Format(@"\\{0}\{1}${2}", hostMachine, path[0], path.Substring(2));
        }

        private static string ConvertToLocal(string path, string hostMachine)
        {
            if (string.IsNullOrEmpty(hostMachine))
                return path;

            path = path.Replace(@"\\{0}\".InvariantCultureFormat(hostMachine), "");
            path = path.Replace("$", ":");

            return path;
        }

        public static void DetachVM(string name)
        {
            DetachVM(name, null);
        }

        public static void DetachVM(string name, string hostMachine)
        {
            var hyperVRootPath = ConvertToUNC(_defaultHyperVRootPath, hostMachine);

            if (!Directory.Exists(hyperVRootPath))
                throw new Exception(string.Format("The Hyper-V root path [{0}] is not valid.", hyperVRootPath));

            var files = (from file in Directory.GetFiles(Path.Combine(hyperVRootPath, "Virtual Machines"), "*.xml")
                         let vmName = GetVMName(file, hostMachine)
                         where vmName.T.Equals(name, StringComparison.InvariantCultureIgnoreCase)
                         select new { vmName, file }).ToList();

            if (files.Count > 1)
                throw new Exception("Cannot detach VM. Two or more VM have the [{0}] name.".InvariantCultureFormat(name));

            if (files.Count == 0)
                throw new Exception("Cannot detach VM. Cannot find [{0}] in [{1}].".InvariantCultureFormat(name, hyperVRootPath));

            var configurationFile = files.First().file;
            var document = files.First().vmName.U;
            var machinePath = hostMachine.Maybe(n => @"\\" + n + "\\", "");

            // Make sure we can detach the VM
            ManagementScope scope = new ManagementScope(machinePath + @"root\virtualization", null);
            ManagementObject vm = Utility.GetTargetComputer(name, scope);

            var state = (VMState)Convert.ToInt32(vm.GetPropertyValue("EnabledState"));

            if(state != VMState.Suspended && state != VMState.Disabled)
                throw new Exception(string.Format("Cannot detach [{0}]. It must be either saved or stopped.", name));

            // Now remove all snapshots and configuration file
            var snapshotsQuery = from element in document.XPathSelectElements(@"/configuration/global_settings/snapshots/list//guid")
                                 select Path.Combine(hyperVRootPath + "\\Snapshots", element.Value + ".xml");

            FormatLogMessage("Removing Configuration symlink [{0}]", configurationFile);

            if(enableWrite)
                File.Delete(configurationFile);

            FormatLogMessage("Removing Snapshot symlink [{0}]", string.Join(", ", snapshotsQuery.ToArray()));
            
            if(enableWrite)
                snapshotsQuery.ForEach(s => File.Delete(s));
        }

        private static void CreateSymbolicLink(string lpSymlinkFileName, string lpTargetFileName, int dwFlags)
        {
#if DEBUG
            FormatLogMessage("mklink \"{0}\" \"{1}\"".InvariantCultureFormat(lpSymlinkFileName, lpTargetFileName));
#endif

            if (enableWrite)
            {
                if (!_CreateSymbolicLink(lpSymlinkFileName, lpTargetFileName, dwFlags))
                {
                    throw new Exception("Failed to create symlink [{0}]".InvariantCultureFormat(lpSymlinkFileName));
                }
            }
        }

        private static string GuessOriginalBasePath(string configurationFile)
        {
            var document = XDocument.Load(configurationFile);

            var searchPaths = new []{
                new { 
                    XPath = "//controller0/drive0/pathname",
                    Filter = @"(?<path>.*?)\\Snapshots\\.*?-",
                    Path = "path"
                },
                new { 
                    XPath = "//controller0/drive0/pathname",
                    Filter = @"(?<path>.*?)\\Virtual Hard Disks\\.*?-",
                    Path = "path"
                },
                new { 
                    XPath = "//savedstate/memlocation",
                    Filter = @"(?<path>.*)\\Virtual Machines\\.*?-",
                    Path = "path"
                },
                new { 
                    XPath = "//controller0/drive0/pathname",
                    Filter = @"(?<path>.*)\\.*?.vhd",
                    Path = "path"
                }
            };


            var q = from xPath in searchPaths
                    let path = document.XPathSelectElement(xPath.XPath).Maybe(v => v.Value)
                    where path.HasValue()
                    let match = Regex.Match(path, xPath.Filter)
                    where match.Success
                    select match.Groups[xPath.Path].Value;

            return q.FirstOrDefault();
        }

        private static void SetAcl(string fileName, IEnumerable<string> aclUserName, bool isSymLink, string hostMachine)
        {
            aclUserName.ForEach(u => SetAcl(fileName, u, isSymLink, hostMachine));
        }

        private static void SetAcl(string fileName, string aclUserName, bool isSymLink, string hostMachine)
        {
            var parameters = string.Format("\"{0}\" /grant \"{1}\":(F) {2}", fileName, aclUserName, isSymLink ? "/L" : "/T");

            ProcessStartInfo info = new ProcessStartInfo("icacls", parameters);
            info.CreateNoWindow = true;
            info.UseShellExecute = false;

#if DEBUG
            FormatLogMessage("icacls " + parameters);
#endif

            if (enableWrite)
            {
                if (fileName.StartsWith("\\"))
                {
                    ConnectionOptions connOptions = new ConnectionOptions();
                    connOptions.Impersonation = ImpersonationLevel.Impersonate;
                    connOptions.EnablePrivileges = true;

                    ManagementScope manScope = new ManagementScope(@"\\{0}\ROOT\CIMV2".InvariantCultureFormat(hostMachine), connOptions);
                    manScope.Connect();

                    ObjectGetOptions objectGetOptions = new ObjectGetOptions();
                    ManagementPath managementPath = new ManagementPath("Win32_Process");
                    ManagementClass processClass = new ManagementClass(manScope, managementPath, objectGetOptions);

                    ManagementBaseObject inParams = processClass.GetMethodParameters("Create");

                    inParams["CommandLine"] = @"icacls " + parameters;
                    ManagementBaseObject outParams = processClass.InvokeMethod("Create", inParams, null);

                    if (Convert.ToInt32(outParams["returnValue"]) != 0)
                    {
                        FormatLogMessage("Failed to set remote permissions to [{0}]".InvariantCultureFormat(fileName));
                    }
                }
                else
                {
                    var p = Process.Start(info);

                    p.WaitForExit();

                    if (p.ExitCode != 0)
                        FormatLogMessage("Failed to set permissions to [{0}]".InvariantCultureFormat(fileName));
                }
            }
        }

        private static void UpdatePaths(string configurationFile, string originalBasePath, string path)
        {
            FormatLogMessage("Updating paths in [{0}] ([{1}] to [{2}])...", configurationFile, originalBasePath, path);

            if (enableWrite)
            {
                var backupFileName = Enumerable.Range(0, 10000)
                                        .Select(i => configurationFile + ".bak{0:000}".InvariantCultureFormat(i))
                                        .Where(n => !File.Exists(n))
                                        .First();

                File.Copy(configurationFile, backupFileName);

                var text = File.ReadAllText(configurationFile).Replace(originalBasePath, path);

                File.WriteAllText(configurationFile, text, Encoding.Unicode);
            }
        }

        private static Tuple<string, XDocument> GetVMName(string configurationFile, string hostMachine)
        {
            configurationFile = GetSymbolicLinkTarget(configurationFile, hostMachine);

            using (var reader = XmlTextReader.Create(File.Open(configurationFile, FileMode.Open, FileAccess.ReadWrite, FileShare.ReadWrite)))
            {
                // The file may be locked by Hyper-V, but we can still open it by using Read/Write sharing
                var document = XDocument.Load(reader);

                return new Tuple<string, XDocument>()
                {
                    T = document.XPathSelectElement("/configuration/properties/name").Value,
                    U = document
                };
            }
        }

        private static string GetSymbolicLinkTarget(string configurationFile, string hostMachine)
        {
            if (configurationFile.StartsWith("\\"))
            {
                // The discovery of the target of a remote symbolic link seems to be problematic.
                // We overcome this by looking up at "dir" output. Not really elegant, but it does the job.
                // 19/03/2009  10:16    <SYMLINK>      1FFC282D-2AC3-4615-9109-01C66CB73E76.xml [C:\hyper-v\Virtual Machines\testvm\Virtual Machines\1FFC282D-2AC3-4615-9109-01C66CB73E76.xml]

                ProcessStartInfo info = new ProcessStartInfo("cmd", "/C dir \"{0}\"".InvariantCultureFormat(configurationFile));
                info.RedirectStandardOutput = true;
                info.UseShellExecute = false;
                info.CreateNoWindow = true;

                var p = Process.Start(info);
                p.WaitForExit();

                var output = p.StandardOutput.ReadToEnd();

                var m = Regex.Match(output, @".*?<SYMLINK>.*?\[(?<target>.*)\]");

                if (m.Success)
                {
                    configurationFile = ConvertToUNC(m.Groups["target"].Value, hostMachine);
                }
            }
            return configurationFile;
        }

        private static void FormatLogMessage(string format, params object[] parms)
        {
            LogMessage(string.Format(format, parms));
        }
    }
}
